import torch
import torch.nn.functional as F

# 假设model是一个已经训练好的语言模型，它接受输入序列并输出下一个词的概率分布

def top_k_sampling(logits, k=5):
    """
    从logits中选择概率最高的k个元素进行采样
    :param logits: 模型输出的logits（未经过softmax的原始输出）
    :param k: 要选择的最高概率元素的数量
    :return: 采样的索引
    """
    probs = F.softmax(logits, dim=-1)  # 应用softmax得到概率分布
    top_k_probs, top_k_indices = probs.topk(k)  # 获取概率最高的k个元素及其索引
    top_k_probs = top_k_probs.squeeze(0)  # 去除批次维度（假设logits是单个样本的输出）
    probs_sum = torch.sum(top_k_probs)
    probs = top_k_probs / probs_sum  # 重新归一化概率
    cumprobs = torch.cumsum(probs, dim=0)  # 计算累积概率
    r = torch.rand(1).item()  # 生成一个[0, 1)之间的随机数
    idx = torch.searchsorted(cumprobs, r, right=True)  # 使用二分查找找到随机数落入的区间
    return top_k_indices[idx].item()

def top_p_sampling(logits, p=0.9):
    """
    从logits中选择累积概率达到或超过p的最小元素集合进行采样
    :param logits: 模型输出的logits（未经过softmax的原始输出）
    :param p: 累积概率的阈值
    :return: 采样的索引
    """
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)  # 对logits进行降序排序
    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)  # 计算累积概率
    cutoff = torch.searchsorted(cumulative_probs, p, right=True)  # 找到累积概率达到或超过p的位置
    if cutoff.item() == 0:
        # 如果所有词的累积概率都小于p，则随机选择一个词
        return torch.randint(0, logits.size(-1), (1,)).item()
    flat_indices = sorted_indices[:cutoff.item()]  # 获取累积概率达到或超过p的词的索引
    probs = F.softmax(logits[flat_indices], dim=-1)  # 重新计算这些词的概率
    probs /= torch.sum(probs)  # 归一化概率
    cumprobs = torch.cumsum(probs, dim=0)  # 计算累积概率
    r = torch.rand(1).item()  # 生成一个[0, 1)之间的随机数
    idx = torch.searchsorted(cumprobs, r, right=True)  # 使用二分查找找到随机数落入的区间
    return flat_indices[idx].item()

# 示例使用
logits = torch.tensor([0.1, 0.4, 0.2, 0.1, 0.2], dtype=torch.float)  # 假设的logits输出

# Top-k 采样
k = 3
sampled_index_k = top_k_sampling(logits, k)
print(f"Top-k 采样 (k={k}): 索引 {sampled_index_k}")

# Top-p 采样
p = 0.7
sampled_index_p = top_p_sampling(logits, p)
print(f"Top-p 采样 (p={p}): 索引 {sampled_index_p}")

